# -*- coding: utf-8 -*-
"""
Stock-GPT Transformer Prediction Script
使用训练好的Transformer模型进行股票预测
"""

import sys
import os
import pandas as pd
import numpy as np
import tensorflow as tf
from tensorflow import keras
import joblib
from datetime import datetime, timedelta
import argparse
# import yfinance as yf  # 不需要，使用本地数据

# 导入自定义模块
from stock_tokenizer import StockTokenizer
from stock_transformer import create_model


def configure_gpu(force_cpu=False):
    """配置GPU设备，支持CPU回退"""
    if force_cpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        print("ℹ️ 强制使用CPU模式")
        
        # 优化CPU多线程配置
        import multiprocessing
        cpu_count = multiprocessing.cpu_count()
        print(f"🔧 检测到 {cpu_count} 个CPU核心")
        
        # 设置TensorFlow并行线程数
        tf.config.threading.set_inter_op_parallelism_threads(cpu_count)
        tf.config.threading.set_intra_op_parallelism_threads(cpu_count)
        
        # 设置环境变量优化
        os.environ['OMP_NUM_THREADS'] = str(cpu_count)
        os.environ['TF_NUM_INTEROP_THREADS'] = str(cpu_count)
        os.environ['TF_NUM_INTRAOP_THREADS'] = str(cpu_count)
        
        print(f"⚡ CPU优化已启用: {cpu_count} 线程并行")
        return False
        
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            return True
        except Exception as e:
            print(f"⚠️ GPU配置失败，回退到CPU: {e}")
            os.environ['CUDA_VISIBLE_DEVICES'] = ''
            
            # CPU回退时也启用多线程优化
            import multiprocessing
            cpu_count = multiprocessing.cpu_count()
            tf.config.threading.set_inter_op_parallelism_threads(cpu_count)
            tf.config.threading.set_intra_op_parallelism_threads(cpu_count)
            os.environ['OMP_NUM_THREADS'] = str(cpu_count)
            print(f"⚡ CPU优化已启用: {cpu_count} 线程并行")
            return False
    return False


class StockTransformerPredictor:
    """Stock Transformer预测器"""
    
    def __init__(self, stock_code: str, frequency: str = 'daily', force_cpu: bool = False):
        self.stock_code = stock_code
        self.frequency = frequency
        
        # 配置GPU
        configure_gpu(force_cpu)
        
        # 文件路径
        self.data_file = f'stock_data_{stock_code}_{frequency}_with_indicators.csv'
        self.tokenizer_file = f'output/scalers/tokenizer_{stock_code}_{frequency}.pkl'
        self.model_file = f'output/models/transformer_{stock_code}_{frequency}.weights.h5'
        self.scaler_file = f'output/scalers/price_scaler_{stock_code}_{frequency}.pkl'
        
        # 检查必要文件是否存在
        self._check_required_files()
        
        # 加载组件
        self.tokenizer = StockTokenizer()
        self.tokenizer.load(self.tokenizer_file)
        self.price_scaler = joblib.load(self.scaler_file)
        
    def _check_required_files(self):
        """检查必需的文件是否存在"""
        required_files = [
            self.tokenizer_file,
            self.model_file,
            self.scaler_file
        ]
        
        missing_files = [f for f in required_files if not os.path.exists(f)]
        
        if missing_files:
            print("❌ 以下必需文件缺失:")
            for f in missing_files:
                print(f"   - {f}")
            print(f"\n请先运行训练脚本: python3 06_train_transformer.py {self.stock_code} -f {self.frequency}")
            sys.exit(1)
            
    def load_model(self):
        """加载训练好的模型"""
        print("🤖 加载Transformer模型...")
        
        # 推断模型参数（从tokenizer和数据推断）
        vocab_size = self.tokenizer.vocab_size
        
        # 使用与训练时相同的参数
        model_params = {
            'd_model': 128,
            'num_heads': 8,
            'num_layers': 4,
            'dff': 256,
            'max_position': 500,
            'dropout_rate': 0.1
        }
        
        # 创建模型
        self.model = create_model(vocab_size, model_params)
        
        # 构建模型
        dummy_input = tf.constant([[1, 2, 3]], dtype=tf.int32)
        _ = self.model(dummy_input, task='price')
        
        # 加载权重
        self.model.load_weights(self.model_file)
        print(f"✅ 模型加载成功: {self.model_file}")
        
    def get_latest_data(self, lookback_days: int = 60):
        """获取最新的市场数据用于预测"""
        print(f"📊 获取最新 {lookback_days} 天的市场数据...")
        
        if os.path.exists(self.data_file):
            # 从已有数据文件读取最新数据
            df = pd.read_csv(self.data_file)
            df = df.tail(lookback_days)
            print(f"   从本地文件获取了最新 {len(df)} 条数据")
        else:
            print("   本地数据文件不存在，建议先运行完整的数据获取流程")
            return None
            
        return df
        
    def predict_next_prices(self, lookback_days: int = 60):
        """预测下一期的开盘价和收盘价"""
        print("🔮 开始预测...")
        
        # 获取最新数据
        latest_data = self.get_latest_data(lookback_days)
        if latest_data is None:
            return None, None
            
        # 对数据进行tokenize
        sequences = self.tokenizer.encode_stock_data(latest_data, self.stock_code)
        
        # 取最后一个序列作为输入
        if not sequences:
            print("❌ 无法生成有效的token序列")
            return None, None
            
        last_sequence = sequences[-1]
        
        # 填充到模型期望的长度
        max_len = 500  # 与训练时保持一致
        if len(last_sequence) > max_len:
            input_sequence = last_sequence[-max_len:]
        else:
            input_sequence = last_sequence + [self.tokenizer.special_tokens['<PAD>']] * (max_len - len(last_sequence))
            
        # 转换为模型输入格式
        input_data = np.array([input_sequence], dtype=np.int32)
        
        # 模型预测
        print("🧠 模型推理中...")
        prediction_scaled = self.model(input_data, task='price', training=False)
        prediction_scaled = prediction_scaled.numpy()
        
        # 反归一化得到真实价格
        prediction_original = self.price_scaler.inverse_transform(prediction_scaled)[0]
        
        predicted_open = prediction_original[0]
        predicted_close = prediction_original[1]
        
        return predicted_open, predicted_close
        
    def display_prediction(self, predicted_open, predicted_close):
        """显示预测结果"""
        if predicted_open is None or predicted_close is None:
            print("❌ 预测失败")
            return
            
        # 获取当前最新价格作为参考
        latest_data = self.get_latest_data(1)
        if latest_data is not None and len(latest_data) > 0:
            current_close = latest_data.iloc[-1]['close']
            
            # 计算涨跌
            open_change = predicted_open - current_close
            close_change = predicted_close - current_close
            
            open_change_pct = (open_change / current_close) * 100
            close_change_pct = (close_change / current_close) * 100
            
            # 判断趋势
            if close_change > 0:
                trend = "看涨 📈"
            elif close_change < 0:
                trend = "看跌 📉"
            else:
                trend = "持平 ➡️"
                
            print("\n" + "="*50)
            print("🔮 Stock-GPT Transformer 预测结果")
            print("="*50)
            print(f"📊 股票代码: {self.stock_code}")
            print(f"⏰ 预测周期: {self.frequency}")
            print(f"📅 预测时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
            print("-"*50)
            print(f"💰 当前收盘价: {current_close:.2f}")
            print(f"🌅 预测开盘价: {predicted_open:.2f}")
            print(f"🌇 预测收盘价: {predicted_close:.2f}")
            print("-"*50)
            print(f"📊 开盘价变化: {open_change:+.2f} ({open_change_pct:+.2f}%)")
            print(f"📊 收盘价变化: {close_change:+.2f} ({close_change_pct:+.2f}%)")
            print(f"📈 预测趋势: {trend}")
            print("="*50)
        else:
            print("\n" + "="*50)
            print("🔮 Stock-GPT Transformer 预测结果")
            print("="*50)
            print(f"📊 股票代码: {self.stock_code}")
            print(f"🌅 预测开盘价: {predicted_open:.2f}")
            print(f"🌇 预测收盘价: {predicted_close:.2f}")
            print("="*50)


def main():
    parser = argparse.ArgumentParser(description='使用Stock-GPT Transformer进行预测')
    parser.add_argument('stock_code', type=str, help='股票代码')
    parser.add_argument('-f', '--frequency', choices=['daily', 'hourly'], 
                       default='daily', help='时间周期')
    parser.add_argument('--force-cpu', action='store_true', help='强制使用CPU模式')
    parser.add_argument('--lookback', type=int, default=60, 
                       help='用于预测的历史数据天数')
    
    args = parser.parse_args()
    
    print("🔮 Stock-GPT Transformer Prediction")
    print("="*40)
    print(f"📊 股票代码: {args.stock_code}")
    print(f"⏰ 时间周期: {args.frequency}")
    print(f"📅 回看天数: {args.lookback}")
    print("="*40)
    
    try:
        # 创建预测器
        predictor = StockTransformerPredictor(
            stock_code=args.stock_code,
            frequency=args.frequency,
            force_cpu=args.force_cpu
        )
        
        # 加载模型
        predictor.load_model()
        
        # 进行预测
        predicted_open, predicted_close = predictor.predict_next_prices(args.lookback)
        
        # 显示结果
        predictor.display_prediction(predicted_open, predicted_close)
        
    except Exception as e:
        print(f"❌ 预测过程出错: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()